#include "gelu_tiling.h"
#include "register/op_def_registry.h"
#include "tiling/platform/platform_ascendc.h"

namespace optiling
{
  const uint32_t BLOCK_SIZE = 32;
  static ge::graphStatus TilingFunc(gert::TilingContext *context)
  {
    // 定义所需变量
    GeluTilingData tiling;
    uint32_t sizeofdatatype;
    uint32_t totalLengthAligned;
    uint32_t totalLength;
    uint32_t ALIGN_NUM;
    uint32_t blockLength;
    uint32_t tileLength;
    uint32_t lasttileLength;
    auto dt = context->GetInputTensor(0)->GetDataType();

    // sizeofdatatype
    if (dt == 1)
    { // float16
      sizeofdatatype = 2;
    }
    else if (dt == 0)
    { // float
      sizeofdatatype = 4;
    }
    // 平台信息获取

    uint32_t ub_block_num = 4;
    if (totalLength >= 1024*1024){
    ub_block_num = 1024*128;
    }
    uint32_t tile_num;

    // totalLength, totalLengthAligned
    totalLength = context->GetInputTensor(0)->GetShapeSize();
    ALIGN_NUM = BLOCK_SIZE / sizeofdatatype;
    if (totalLength % ALIGN_NUM != 0)
    {
      totalLengthAligned = ((totalLength + ALIGN_NUM - 1) / ALIGN_NUM) * ALIGN_NUM;
    }
    else
    {
      totalLengthAligned = totalLength;
    }

    context->SetBlockDim(1);
    auto block_dim = context->GetBlockDim();
    blockLength = totalLengthAligned / block_dim;
    tile_num = blockLength / ALIGN_NUM / ub_block_num;
    if ((totalLengthAligned / block_dim / ALIGN_NUM) % ub_block_num == 0 ||
        tile_num == 0)
    {
      if (tile_num == 0)
      {
        tile_num = 1;
      }
      if (blockLength < ub_block_num)
      {
        tileLength = ((blockLength / ALIGN_NUM) + 1) / 2 * 2 * ALIGN_NUM;
        lasttileLength = tileLength;
      }
      else
      {
        tileLength = ub_block_num * ALIGN_NUM;
        lasttileLength = tileLength;
      }
    }
    else
    {
      tile_num = tile_num + 1;
      tileLength = ub_block_num * ALIGN_NUM;
      lasttileLength = blockLength - (tile_num - 1) * tileLength;
    }
    tiling.set_blockLength(blockLength);
    tiling.set_tileNum(tile_num);
    tiling.set_tileLength(tileLength);
    tiling.set_lasttileLength(lasttileLength);
    tiling.SaveToBuffer(context->GetRawTilingData()->GetData(),
                        context->GetRawTilingData()->GetCapacity());
    context->GetRawTilingData()->SetDataSize(tiling.GetDataSize());
    size_t *currentWorkspace = context->GetWorkspaceSizes(1);
    currentWorkspace[0] = 0;
    return ge::GRAPH_SUCCESS;
  }
} // namespace optiling

namespace ge
{
  static ge::graphStatus InferShape(gert::InferShapeContext *context)
  {
    const gert::Shape *x1_shape = context->GetInputShape(0);
    gert::Shape *y_shape = context->GetOutputShape(0);
    *y_shape = *x1_shape;
    return GRAPH_SUCCESS;
  }
}

namespace ops
{
  class Gelu : public OpDef
  {
  public:
    explicit Gelu(const char *name) : OpDef(name)
    {
      this->Input("x")
          .ParamType(REQUIRED)
          .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
          .Format({ge::FORMAT_ND, ge::FORMAT_ND})
          .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});
      this->Output("y")
          .ParamType(REQUIRED)
          .DataType({ge::DT_FLOAT16, ge::DT_FLOAT})
          .Format({ge::FORMAT_ND, ge::FORMAT_ND})
          .UnknownShapeFormat({ge::FORMAT_ND, ge::FORMAT_ND});

      this->SetInferShape(ge::InferShape);

      this->AICore()
          .SetTiling(optiling::TilingFunc);
      this->AICore().AddConfig("ascend310b");
    }
  };

  OP_ADD(Gelu);
}

